#include "pointfit.h"

void PointFitter::fit(vector<Point3f> pnts0, vector<Point3f> pnts1, Mat &R, Mat &T)
{
    assert( pnts0.size() == pnts1.size() );
    // step 1
    Point3f center0, center1;
    center0 = Centeroid( pnts0 ); // eq. 6
    center1 = Centeroid( pnts1 ); // eq. 4

    vector<Point3f> q0s;
    vector<Point3f> q1s;
    Minus( pnts0, center0, q0s ); // eq. 7
    Minus( pnts1, center1, q1s ); // eq. 8

    // step 2
    Mat H;
    HMatrix( q0s, q1s, H );       // eq. 11

    // step 3 and 4
    FindRotation( H, R );         // eq. 12 & 13

    T = PtoM( center1 ) - R * PtoM( center0 ); // eq. 10
}

/* Get the centroid of a group of points */
Point3f PointFitter::Centeroid( vector<Point3f> pnts )
{
    Point3f res( 0, 0, 0 ) ;
    for( int i = 0; i < pnts.size(); i ++ )
    {
        res += pnts.at( i );
    }
    res.x /= pnts.size();
    res.y /= pnts.size();
    res.z /= pnts.size();
    return res;
}

/* Translate the given point vector inverse with p */
void PointFitter::Minus(vector<Point3f> pnts, Point3f p, vector<Point3f> &output)
{
    output.clear();
    for( int i = 0; i < pnts.size(); i ++ )
    {
        output.push_back( pnts.at( i ) - p );
    }
}

/* Create the cv::Mat format for a point */
Mat PointFitter::PtoM( Point3f p )
{

    Mat res( 3, 1, CV_64F, Scalar::all(0) );
    res.at<double>(0,0) = p.x;
    res.at<double>(1,0) = p.y;
    res.at<double>(2,0) = p.z;
    return res;
}

/* Calculate the H matrix ( the number 11 equation on the paper) */
void PointFitter::HMatrix( vector<Point3f> q0s, vector<Point3f> q1s, Mat& res )
{
    res = Mat::zeros( 3, 3, CV_64FC1 );
    for( int i = 0; i < q0s.size(); i ++ )
    {
        Mat q0 = PtoM( q0s.at( i ) );
        Mat q1 = PtoM( q1s.at( i ) );
        Mat t;
        transpose( q1, q1 );
        t = q0 * q1;
        res += t;
    }
}

/* Calculate the rotation from SVD */
void PointFitter::FindRotation( Mat H, Mat& R )
{
    Mat w, u, vt;

    SVD::compute( H, w, u, vt );
    transpose( vt, vt );
    transpose( u, u );

    R = vt * u;
}
